#!/usr/bin/env python
# -*- coding:utf-8 -*-
"""
@File   : ssh_tools.py
@Time : 2024/12/9 11:19
@Author: Wangwc
@Describe:  ……
"""
import os
import shutil
import sys
import paramiko
import select
import time

from commons.logger import log
from utils.requests_tools import retry_func


def get_resource_path(filename):
    """获取资源文件的路径"""
    try:
        log.info(f"开始获取资源文件路径：{filename}")

        # 判断程序是否被打包成可执行文件
        if getattr(sys, 'frozen', False):
            # 如果是被打包成exe的
            exe_directory = os.path.dirname(sys.executable)
            log.debug(f"程序被打包为可执行文件，目录：{exe_directory}")
        else:
            # 直接运行脚本
            exe_directory = os.path.dirname(__file__)
            log.debug(f"直接运行脚本，目录：{exe_directory}")

        # 构建完整的文件路径
        config_file_path = os.path.join(exe_directory, filename)
        log.info(f"资源文件完整路径：{config_file_path}")

        return config_file_path

    except Exception as e:
        log.error(f"获取资源文件路径时发生错误：{e}", exc_info=True)
        raise


def convert_windows_to_linux(file_path):
    """将Windows系统字符转换成Linux系统字符"""
    log.info(f"开始转换文件: {file_path}")

    if not os.path.exists(file_path):
        log.error(f"文件不存在: {file_path}")
        raise FileNotFoundError(f"文件不存在: {file_path}")

    try:
        with open(file_path, 'r', newline='', encoding='utf-8') as infile:
            log.debug(f"打开文件进行读取: {file_path}")
            content = infile.read()
            log.debug("读取文件内容成功")

        # 替换 CRLF 换行符为 LF
        linux_content = content.replace('\r\n', '\n')
        log.debug("完成换行符替换")

        # 覆盖原文件
        with open(file_path, 'w', newline='', encoding='utf-8') as outfile:
            log.debug(f"打开文件进行写入: {file_path}")
            outfile.write(linux_content)
            log.info(f"成功转换并覆盖文件: {file_path}")

    except Exception as e:
        log.error(f"转换文件时发生错误: {e}")
        raise

    finally:
        log.debug("文件转换操作结束")


def create_ssh_client(hostname, username, password,port=22):
    """创建SSH客户端并记录登录尝试"""
    log.info(f"尝试连接到主机 {hostname} 使用用户 {username}")

    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        log.debug("设置SSH客户端的缺失主机密钥策略为AutoAddPolicy")
        ssh_client.connect(
            hostname=hostname,
            port=port,
            username=username,
            password=password,
            timeout=60,
            auth_timeout=60
        )
        log.info(f"成功登录到主机 {hostname} 使用用户 {username}")
        return ssh_client

    except paramiko.AuthenticationException:
        log.error(f"用户 {username}登录主机 {hostname}失败：认证错误。 ")
        raise

    except paramiko.SSHException as e:
        log.error(f"SSH连接失败：{e}. 主机: {hostname} 用户: {username}")
        raise

    except Exception as e:
        log.error(f"连接失败：{e}. 主机: {hostname} 用户: {username}")
        raise

    finally:
        log.debug("SSH客户端连接尝试结束")


@retry_func()
def create_ssh_client_as_root(hostname: str, username: str, password: str,port=22):
    """
    创建SSH客户端并记录登录尝试
    :param hostname: 主机名
    :param username: root用户名
    :param password: root密码
    """

    log.info(f"尝试连接到主机 {hostname} 使用用户 {username}")

    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        ssh_client.connect(
            hostname=hostname,
            port=port,
            username=username,
            password=password,
            timeout=60,
            auth_timeout=60
        )
        log.info(f"成功登录到主机 {hostname} 使用用户 {username}")

        # 使用 exec_command 获取当前用户
        stdin, stdout, stderr = ssh_client.exec_command('whoami')
        current_user = stdout.read().decode().strip()

        if current_user != 'root':
            log.info(f"当前用户 {current_user} 不是 root，尝试通过合法交互方式切换到 root 用户")
            shell = ssh_client.invoke_shell()
            if not switch_to_root(shell, password):
                raise Exception("切换到 root 用户失败")

            modify_ssh_config(shell, password)
            restart_ssh_service(shell)

            # 重新使用 root 用户登录
            ssh_client.close()
            ssh_client = reconnect_as_root(hostname, password)

        return ssh_client

    except paramiko.AuthenticationException:
        log.error(f"用户 {username} 登录主机 {hostname} 失败：认证错误。")
        raise

    except paramiko.SSHException as e:
        log.error(f"SSH连接失败：{e}. 主机: {hostname} 用户: {username}")
        raise

    except Exception as e:
        log.error(f"连接失败：{e}. 主机: {hostname} 用户: {username}")
        raise

    finally:
        log.debug("SSH客户端连接尝试结束")


def execute_command(shell, command: str, timeout=5):
    """执行命令并设置超时机制"""
    shell.send(command + '\n')
    ready, _, _ = select.select([shell], [], [], timeout)
    if ready:
        output = shell.recv(1024).decode()
        if "command not found" in output:
            log.error(f"命令 {command} 在服务器上不存在或不可执行")
            raise Exception("命令执行失败，可能命令不存在")
        return output
    else:
        log.error(f"命令 {command} 执行超时")
        return ""


def switch_to_root(shell, password: str) -> bool:
    """通过sudo切换到root用户"""
    shell.send('sudo -i\n')
    time.sleep(1)
    output = shell.recv(1024).decode()
    if "[sudo] password for" in output:
        shell.send(f"{password}\n")
        time.sleep(1)
        new_output = shell.recv(1024).decode()
        if "Sorry, try again" in new_output or "Permission denied" in new_output:
            log.error("切换到 root 用户失败：密码错误或权限不足")
            return False
        log.info("成功切换到 root 用户")
        return True
    log.error("切换到 root 用户失败：未提示输入密码")
    return False


def modify_ssh_config(shell, password: str):
    """修改 /etc/ssh/sshd_config 相关配置"""
    config_settings = {
        "PermitRootLogin": "yes",
        "ClientAliveInterval": "60",
        "ClientAliveCountMax": "3",
        "UseDNS": "no"
    }

    # 尝试修改配置的重试机制
    log.info("开始修改 SSH 配置文件，设置root用户可以登录ssh，并关闭DNS解析")
    try:
        for key, value in config_settings.items():
            check_command = f"grep -i '{key}' /etc/ssh/sshd_config | grep -v '^#'"
            output = execute_command(shell, check_command)
            current_value = output.strip().split()[-1] if output.strip() else ""
            if current_value != value:
                command = f"sed -i 's/^{key}.*$/{key} {value}/' /etc/ssh/sshd_config"
                execute_command(shell, command)

        # 修改root密码
        shell.send("passwd root\n")
        time.sleep(1)
        output = shell.recv(4096).decode()
        if "New password:" in output:
            shell.send(f"{password}\n")
            time.sleep(1)
            if "Retype new password:" in shell.recv(1024).decode():
                shell.send(f"{password}\n")
                time.sleep(1)
                output = shell.recv(1024).decode()  # 读取确认密码提示
                if "passwd: password updated successfully" in output:
                    log.info("成功修改 root 密码")
                    return  # 成功后退出函数
                else:
                    log.warning("修改 root 密码失败")
                    raise Exception("修改 root 密码失败")
            else:
                log.warning("重新输入密码提示未出现，修改 root 密码失败")
                raise Exception("修改 root 密码失败")
        else:
            log.warning("输入新密码提示未出现，修改 root 密码失败")
            raise Exception("修改 root 密码失败")

    except Exception as e:
        log.error("修改 SSH 配置和密码失败，已达到最大重试次数")
        raise Exception("修改 SSH 配置和密码失败，已达到最大重试次数")


def restart_ssh_service(shell):
    """重新启动 SSH 服务"""
    log.info("开始重新启动 SSH 服务")
    shell.send('sudo systemctl restart sshd\n')
    time.sleep(3)
    output = shell.recv(1024).decode()
    log.info("SSH 服务已重新启动")


def reconnect_as_root(hostname: str, password: str):
    """以root用户重新连接"""
    log.info("尝试以 root 用户重新连接")
    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    for _ in range(3):  # 尝试多次登录
        try:
            ssh_client.connect(
                hostname=hostname,
                port=22,
                username='root',
                password=password,
                timeout=60,
                auth_timeout=60
            )
            log.info("成功以 root 用户身份重新登录")
            return ssh_client
        except paramiko.AuthenticationException:
            log.error("以 root 用户重新登录失败，等待后再次尝试...")
            time.sleep(2)
    log.error("多次尝试后，以 root 用户重新登录仍失败")
    raise Exception("以 root 用户重新登录失败")


def sftp_upload_files(ssh_client: paramiko.SSHClient, file_path: str, remote_path: str):
    """上传文件到服务器"""
    log.info(f"开始上传文件: {file_path} 到远程路径: {remote_path}")

    # 获取文件的绝对路径
    absolute_file_path = os.path.abspath(file_path)

    if not os.path.exists(absolute_file_path):
        log.error(f"文件不存在: {absolute_file_path}")
        raise FileNotFoundError(f"文件不存在: {absolute_file_path}")

    # 检查文件权限
    if not os.access(absolute_file_path, os.R_OK):
        log.error(f"没有权限读取文件: {absolute_file_path}")
        raise PermissionError(f"没有权限读取文件: {absolute_file_path}")

    try:
        # 将 Windows 换行符转换为 Linux 格式
        log.debug(f"转换文件换行符: {absolute_file_path}")
        convert_windows_to_linux(absolute_file_path)

        # 使用SFTP客户端
        log.debug("打开SFTP传输通道")
        with ssh_client.open_sftp() as sftp:
            # 上传文件
            log.debug(f"上传文件: {absolute_file_path} 到远程路径: {remote_path}")
            sftp.put(absolute_file_path, remote_path, confirm=True)
            log.info(f"成功上传文件: {absolute_file_path} 到远程路径: {remote_path}")

    except FileNotFoundError as fnf_error:
        log.error(f"文件未找到错误: {fnf_error}")
        raise

    except PermissionError as perm_error:
        log.error(f"权限错误: {perm_error}")
        raise

    except paramiko.SSHException as ssh_error:
        log.error(f"SFTP传输失败: {ssh_error}")
        raise

    except Exception as e:
        log.error(f"发生未知错误: {e}")
        raise

    finally:
        log.debug("文件上传操作结束")


def download_file(ssh_client, remote_log_path, local_log_path):
    """下载日志文件"""
    try:
        log.info(f"开始下载日志文件：{remote_log_path} 到本地路径：{local_log_path}")

        # 创建SFTP客户端
        sftp_client = ssh_client.open_sftp()
        log.debug("SFTP客户端已创建")

        # 如果本地文件已存在，则删除
        if os.path.exists(local_log_path):
            log.info(f"本地文件{local_log_path}已存在，删除旧文件")
            os.remove(local_log_path)

        # 尝试获取远程文件的属性并下载
        try:
            sftp_client.stat(remote_log_path)
            log.debug(f"远程文件{remote_log_path}存在，准备下载")
            sftp_client.get(remote_log_path, local_log_path)
            log.info(f"成功下载文件到{local_log_path}")

        except FileNotFoundError:
            log.error(f"远程文件{remote_log_path}不存在")
            raise FileNotFoundError(f"远程文件{remote_log_path}不存在")
        finally:
            # 关闭SFTP客户端
            sftp_client.close()
            log.debug("SFTP客户端已关闭")
    except paramiko.SSHException as ssh_error:
        log.error(f"SSH连接错误：{ssh_error}", exc_info=True)
        raise

    except IOError as io_error:
        log.error(f"文件操作错误：{io_error}", exc_info=True)
        raise

    except Exception as e:
        log.error(f"发生未知错误：{e}", exc_info=True)
        raise


def check_process_activity(ssh_client, pid):
    """检查指定PID的进程是否存在"""
    command = f'ps -p {pid} -o tty=,time=,cmd='
    try:
        log.info(f'开始检查PID为{pid}的进程活动')
        stdin, stdout, stderr = ssh_client.exec_command(command)
        log.debug(f'执行命令：{command}')

        output = stdout.read().decode().strip()
        error_output = stderr.read().decode().strip()

        if error_output:
            log.error(f"检查进程时发生错误：{error_output}")
            return None

        if output:
            log.info(f"进程活动信息：{output}")
        else:
            log.info(f"未找到PID为{pid}的进程")

        return output

    except paramiko.SSHException as ssh_error:
        log.error(f"SSH命令执行失败：{ssh_error}", exc_info=True)
        raise

    except Exception as e:
        log.error(f"发生错误：{e}", exc_info=True)
        raise

    finally:
        log.info("检查进程活动操作结束")


@retry_func()
def create_raid_with_shell2(ssh_client, raid_type=1, disk="sdb sdc", disk_size=20):
    """
    通过已经建立的SSH连接，使用invoke_shell方式执行创建RAID的操作

    :param ssh_client: 已经建立好的SSH连接客户端对象
    :param raid_type: RAID类型，对应输入的数字（如1代表raid1等）
    :param disk: 需要创建RAID的磁盘列表，格式为字符串，多个磁盘用空格隔开（如 "sda sdb"）
    :param disk_size: 需要创建的磁盘大小，单位为GB，输入整数（直接回车则使用剩余大小，这里按传入整数处理）
    :return: 如果执行成功返回True，出现异常返回False
    """
    log.info("开始执行创建RAID操作")
    try:
        shell = ssh_client.invoke_shell()
        log.info("已获取SSH交互式Shell通道，等待通道准备好...")
        time.sleep(1)  # 等待一会儿，确保通道准备好
        log.info("通道准备就绪，发送启动创建RAID脚本的命令")

        # 发送命令启动创建RAID的脚本
        shell.send('./create_raid.sh\n')
        time.sleep(1)
        log.info("已发送启动创建RAID脚本的命令，等待1秒")

        # 依次发送交互式输入的内容
        shell.send(f"{raid_type}\n")
        log.info(f"已发送RAID类型输入内容: {raid_type}，等待1秒")
        time.sleep(1)
        shell.send(f"{disk}\n")
        log.info(f"已发送磁盘列表输入内容: {disk}，等待1秒")
        time.sleep(1)
        shell.send(f"{disk_size}\n")
        log.info(f"已发送磁盘大小输入内容: {disk_size}，等待1秒")
        time.sleep(1)

        # 等待一段时间让创建RAID脚本执行完成，这里等待30秒，可根据实际情况调整
        log.info("已发送所有交互式输入内容，开始等待创建RAID脚本执行完成，等待时间20秒")
        time.sleep(20)

        # 获取命令执行结果
        output = shell.recv(65535).decode('utf-8')
        if "阵列创建和数据库插入完成。" in output and "重启存储服务完成！" in output:
            log.info("RAID创建成功，输出信息如下:\n" + output)
            return True
        else:
            log.error(f"执行创建RAID脚本出现错误，输出信息如下:\n{output}")
            return False
    except Exception as e:
        log.error(f"创建RAID过程出现异常: {e}")
        return False
    finally:
        log.info("创建RAID操作执行结束")

def get_disk_sizes(ssh_client, disks):
    """
    通过 SSH 连接获取指定硬盘的大小
    :param ssh_client: 已经建立好的 SSH 连接客户端对象
    :param disks: 硬盘列表
    :return: 硬盘及其大小的字典
    """
    disk_sizes = {}
    for disk in disks:
        stdin, stdout, stderr = ssh_client.exec_command(f"blockdev --getsize64 /dev/{disk}")
        output = stdout.read().decode('utf-8').strip()
        if output:
            size = int(output) / (1024 ** 3)  # 转换为 GB
            disk_sizes[disk] = size
    return disk_sizes

def create_raid_with_shell(ssh_client, raid_type=1, min_disk_size=100, disk_size=20):
    """
    通过已经建立的 SSH 连接，使用 invoke_shell 方式执行创建 RAID 的操作

    :param ssh_client: 已经建立好的 SSH 连接客户端对象
    :param raid_type: RAID 类型，对应输入的数字（如 1 代表 raid1 等）
    :param min_disk_size: 硬盘最小大小，单位为 GB，用于筛选合适的硬盘
    :param disk_size: 需要创建的磁盘大小，单位为 GB，输入整数（直接回车则使用剩余大小，这里按传入整数处理）
    :return: 如果执行成功返回 True，出现异常返回 False
    """
    log.info("开始执行创建 RAID 操作")
    try:
        # 检查硬盘大小
        disks = ['sda', 'sdb', 'sdc', 'sdd']
        disk_sizes = get_disk_sizes(ssh_client, disks)
        suitable_disks = [disk for disk, size in disk_sizes.items() if size >= min_disk_size]

        if len(suitable_disks) < 2:
            log.error(f"可用的大于 {min_disk_size}GB 的硬盘数量不足，无法创建 RAID 1 存储池。")
            return False

        # 选择前两个合适的硬盘
        selected_disks = " ".join(suitable_disks[:2])

        shell = ssh_client.invoke_shell()
        log.info("已获取 SSH 交互式 Shell 通道，等待通道准备好...")
        time.sleep(1)  # 等待一会儿，确保通道准备好
        log.info("通道准备就绪，发送启动创建 RAID 脚本的命令")

        # 发送命令启动创建 RAID 的脚本
        shell.send('./create_raid.sh\n')
        time.sleep(1)
        log.info("已发送启动创建 RAID 脚本的命令，等待 1 秒")

        # 依次发送交互式输入的内容
        shell.send(f"{raid_type}\n")
        log.info(f"已发送 RAID 类型输入内容: {raid_type}，等待 1 秒")
        time.sleep(1)
        shell.send(f"{selected_disks}\n")
        log.info(f"已发送磁盘列表输入内容: {selected_disks}，等待 1 秒")
        time.sleep(1)
        shell.send(f"{disk_size}\n")
        log.info(f"已发送磁盘大小输入内容: {disk_size}，等待 1 秒")
        time.sleep(1)

        # 等待一段时间让创建 RAID 脚本执行完成，这里等待 30 秒，可根据实际情况调整
        log.info("已发送所有交互式输入内容，开始等待创建 RAID 脚本执行完成，等待时间 20 秒")
        time.sleep(20)

        # 获取命令执行结果
        output = shell.recv(65535).decode('utf-8')
        if "阵列创建和数据库插入完成。" in output and "重启存储服务完成！" in output:
            log.info("RAID 创建成功，输出信息如下:\n" + output)
            return True
        else:
            log.error(f"执行创建 RAID 脚本出现错误，输出信息如下:\n{output}")
            return False
    except Exception as e:
        log.error(f"创建 RAID 过程出现异常: {e}")
        return False
    finally:
        log.info("创建 RAID 操作执行结束")


if __name__ == "__main__":
    hostname = "172.17.21.129"
    username = "test3"
    password = "Aa123456@"
    ssh_client = create_ssh_client_as_root(hostname, username, password)
    create_raid_with_shell(ssh_client, raid_type=1,disk_size=10)
